package com.adam.example.aggregation.bucket;

import co.elastic.clients.elasticsearch._types.DistanceUnit;
import co.elastic.clients.elasticsearch._types.aggregations.AggregationRange;
import co.elastic.clients.elasticsearch.core.SearchRequest;
import com.adam.example.util.CollectionStreamUtil;
import com.adam.example.util.JsonUtil;
import com.adam.example.util.OkHttpUtil;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;

import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.Map;

public class OkHttpGeoDistanceAggregationExample {

    public static void main(String[] args) throws IOException {
        SearchRequest searchRequest = SearchRequest.of(builder -> builder
                        .index("test_geo_point")
                        .size(0)
                        .aggregations("agg", aggregationBuilder -> aggregationBuilder
                                .geoDistance(geoDistanceAggregationBuilder -> geoDistanceAggregationBuilder
                                        .field("location")
                                        .origin(geoLocationBuilder -> geoLocationBuilder
                                                .text("34.3412700000,108.9398400000")
                                        )
                                        .unit(DistanceUnit.Kilometers)
                                        .ranges(Arrays.asList(
                                                AggregationRange.of(aggregationRangeBuilder -> aggregationRangeBuilder
                                                        .to("500")
                                                ),
                                                AggregationRange.of(aggregationRangeBuilder -> aggregationRangeBuilder
                                                        .from("500").to("1000")
                                                ),
                                                AggregationRange.of(aggregationRangeBuilder -> aggregationRangeBuilder
                                                        .from("1000")
                                                )
                                        ))
                                )
                        )
        );
        String requestJson = JsonUtil.searchRequestToJson(searchRequest);
//        System.out.println(requestJson);
        ObjectMapper objectMapper = new ObjectMapper();
        Map<String, Object> map = objectMapper.readValue(requestJson, new TypeReference<Map<String, Object>>() {});
        List<Map<String, Object>> rangesList = CollectionStreamUtil.build(map)
                .getKey("aggregations").getKey("agg").getKey("geo_distance").getKey("ranges").getListMap();
        for(Map<String, Object> map2: rangesList) {
            String fromString = (String) map2.get("from"),
                    toString = (String) map2.get("to");
            if(fromString != null) {
                long from = Long.parseLong(fromString);
                map2.put("from", from);
            }
            if(toString != null) {
                long to = Long.parseLong(toString);
                map2.put("to", to);
            }
        }
        String processedRequestJson = objectMapper.writeValueAsString(map);
        System.out.println("处理过的请求json:" + processedRequestJson);

        String responseJson = OkHttpUtil.postSearch("test_geo_point", processedRequestJson);
        System.out.println("响应json:" + responseJson);
        Map<String, Object> rMap = objectMapper.readValue(responseJson, new TypeReference<Map<String, Object>>() {});
        List<Map<String, Object>> bucketList = CollectionStreamUtil.build(rMap)
                        .getKey("aggregations").getKey("agg").getKey("buckets").getListMap();
        for(Map<String, Object> bucketMap : bucketList) {
            System.out.println(bucketMap.get("key") + ":" + bucketMap.get("doc_count"));
        }
    }

}
